from random import shuffle


#MARK: constants

_TAGS_PATH			= "netflix_tags/tag_info.txt"
_RATING_PATH_FOLDER	= "netflix_data/netflix/download/training_set"
_RATING_PATH_BASE	= _RATING_PATH_FOLDER + "/mv_%07d.txt"

_DEGREE_TAGS_MAPPED = ["_RATING"]

_DEGREE_TAGS = {"PG"	:	(_DEGREE_TAGS_MAPPED[0], 0.0),
				"PG-13"	:	(_DEGREE_TAGS_MAPPED[0], 0.5),
				"R" 	:	(_DEGREE_TAGS_MAPPED[0], 1.0)}


#MARK: movies

def read_movie_features(movie_max, tag_max):
	# pass 1: read tag counts
	tag_counts = {}
	
	f = open(_TAGS_PATH)
	
	for line in iter(f):
		fields = line.split("\t")
		assert len(fields) == 5
		
		tags = fields[-1].split(";")
		
		for tag_dirty in tags:
			tag = tag_dirty.strip()
			if not len(tag): continue
			
			tag, val = fix_degree_tag(tag)
			
			if not tag in tag_counts: tag_counts[tag] = 0
			tag_counts[tag] += 1
	
	f.close()
	
	# layout feature vector
	tag_vector = tag_counts.items()
	tag_vector.sort(key=lambda pair: pair[1], reverse=True)
	
	if tag_max != None:
		tag_vector = tag_vector[:min(tag_max, len(tag_vector))]
	
	# create map from tag to index into feature vector
	index_map = {}
	
	for i, (tag, count) in enumerate(tag_vector):
		index_map[tag] = i
	
	# pass 2: read movie features
	movies = {}
	miss_count = 0
	
	f = open(_TAGS_PATH)
	
	for line in iter(f):
		fields = line.split("\t")
		assert len(fields) == 5
		
		movie_id = int(fields[0])
		tags = fields[-1].split(";")
		
		# populate features
		features = [0 for i in range(len(tag_vector))]
		
		for tag_dirty in tags:
			tag = tag_dirty.strip()
			if not len(tag): continue
			
			tag, val = fix_degree_tag(tag)
			
			if tag in index_map:
				features[index_map[tag]] = val
		
		# skip if none
		if sum(features) == 0:
			miss_count += 1; continue;
		
		movies[movie_id] = features
		
		# limit movies if requested
		if movie_max != None and len(movies) > movie_max:
			break
	
	f.close()
	
	return movies, tag_vector, miss_count

def is_tag_visible(tag):
	return tag not in _DEGREE_TAGS_MAPPED

def fix_degree_tag(orig_tag):
	if orig_tag in _DEGREE_TAGS:	return _DEGREE_TAGS[orig_tag]
	else:							return orig_tag, 1


#MARK: user ratings

def read_user_ratings(user_ids, movies):
	ratings = {}
	for user_id in user_ids:
		ratings[user_id] = []
	
	for movie_id, features in movies.iteritems():
		try:		f = open(_RATING_PATH_BASE % movie_id)
		except:		raise "data missing for movie %d" % movie_id
		
		for line in iter(f):
			fields = line.split(",")
			if len(fields) != 3: continue
			
			user_id = int(fields[0])
			rating  = int(fields[1])
			
			if user_id in ratings:
				ratings[user_id].append((movie_id, rating))
		
		f.close()
	
	return ratings


#MARK: users

def read_users(movies, user_max):
	user_ids = set()
	
	# collect all users who have rated at least one movie
	for movie_id, _junk in movies.iteritems():
		f = open(_RATING_PATH_BASE % movie_id)
		
		for line in iter(f):
			fields = line.split(",")
			if len(fields) == 3:
				user_ids.add(int(fields[0]))
		
		f.close()
	
	# make randomly-ordered list
	user_ids = list(user_ids)
	shuffle(user_ids)
	
	if user_max != None:
		user_ids = user_ids[:user_max]
	
	return user_ids
